#include"../util.h"
#include"config.h"
#include"edge_checker\raw_distance_checker.h"
#include"graph\graph.h"
#include"default_model.h"
#include"../../log/log_mgr.h"

DefaultModel::DefaultModel(DistCalc *dist_calc, LanguageModel *lang_model, TranslateModel *trans_model)
{
	this->dist_calc = dist_calc;
	this->lang_model = lang_model;
	this->trans_model = trans_model;
}

DefaultModel::DefaultModel(DistCalc *dist_calc, LanguageModel *lang_model, TranslateModel *trans_model, AbbrModel *abbr_model)
{
	this->dist_calc = dist_calc;
	this->lang_model = lang_model;
	this->trans_model = trans_model;
	this->abbr_model = abbr_model;
}

DefaultModel::DefaultModel(DistCalc *dist_calc, LanguageModel *lang_model, TranslateModel *trans_model, AbbrModel *abbr_model, CorrectionModel *corr_model)
{
	this->dist_calc = dist_calc;
	this->lang_model = lang_model;
	this->trans_model = trans_model;
	this->abbr_model = abbr_model;
	this->corr_model = corr_model;
}

void DefaultModel::co_occur(vector<string> &term_set, vector<string> &term_list, int cur, int start, int freq, map<string, int> &lm)
{
	if(cur == atoi(Config::get("gram").c_str()) + 1)return;
	for(int i = start; i < term_list.size(); ++ i)
	{
		term_set.push_back(term_list[i]);
		vector<string> term_set_s = util::sort_v(term_set);
		string s;
		for(int j = 0; j < term_set_s.size(); ++ j)
		{
			s += term_set_s[j];
			s += " ";
		}
		map<string, int>::iterator p2f_itr = lm.find(s);
		if(p2f_itr == lm.end())
			lm.insert(pair<string, int>(s, freq));
		else
			(*p2f_itr).second += freq;
		co_occur(term_set, term_list, cur + 1, i + 1, freq, lm);
		term_set.pop_back();
	}
}

void DefaultModel::co_occur(string query, int freq, map<string, int> &lm)
{
	vector<string> term_set;
	vector<string> term_list = util::tokenize(query, " ");		
	co_occur(term_set, term_list, 1, 0, freq, lm);
}

void DefaultModel::raw_grouping(map<string, int> &q2f, vector<vector<string>> &gro, DistCalc *dc)
{
	map<string, int>::iterator q_itr;
	QueryGraph graph;
	for(q_itr = q2f.begin(); q_itr != q2f.end(); ++ q_itr)
		graph.add_node((*q_itr).first);
	graph.update_edge(RawDistanceChecker(dc));
	//cout << "graph.update_edge done\n";
	graph.flood(gro);
	//cout << "graph.flood done\n";
}

string DefaultModel::raw_refine(string str, map<string, int> &lm, map<string, string> &w2c)
{
	string label;
	vector<string> terms = util::tokenize(str, " ");
	for(int j = 0 ; j < terms.size(); ++ j){
		string word = terms[j] + " ";
		string cand, c_cand;
		map<string, string>::iterator w2c_itr = w2c.find(word);
		if(w2c_itr == w2c.end()){
			if(word.length() < 3)cand = word;
			else{
				map<string, pair<double, double>> can_info;
				map<string, pair<double, double>>::iterator can_itr;
				map<string, int>::iterator itr;
				double score = NEG_INF, c_score = 0, freq, total_freq = 0;
				set<string> neighbour;
				for(itr = lm.begin(); itr != lm.end(); ++ itr){
					c_cand = (*itr).first;
					freq = (*itr).second;
					c_score = dist_calc->dist(word, c_cand);
					if(!c_cand[0] == word[0])continue;
					if(c_score > atof(Config::get("hard_term_thr").c_str()))
						continue;
					c_score = c_score/(word.length());
					if(c_score > atof(Config::get("prob_term_thr").c_str()))
						continue;
					can_info.insert(pair<string, pair<double, double>>(c_cand,
									pair<double, double>(c_score, freq)));
					total_freq += freq;
				}
				for(can_itr = can_info.begin(); can_itr != can_info.end(); ++ can_itr){
					c_cand = (*can_itr).first;
					c_score = (*can_itr).second.first;
					freq = (*can_itr).second.second;
					c_score = exp(-c_score)*freq/total_freq;
					if(c_score > score){
						score = c_score;
						cand = c_cand;
					}
				}
			}
			if(cand == "")cand = word;
			w2c.insert(pair<string, string>(word, cand));
		}
		else
			cand = (*w2c_itr).second;					
		label += cand;
	}
	return label;
}

void DefaultModel::elect_label(vector<vector<string>> &gro, map<string, int> &q2f, map<string, int> &lm, map<string, vector<pair<string, int>>> &l2g)
{
	set<vector<string>>::iterator itr;
	map<string, string> word2corr;
	for(int i = 0; i < gro.size(); ++ i)
	{
		vector<string> elem = gro[i];
		vector<pair<string, int>> elem_info;
		string label, t_label;
		double score = 0, t_score = 0;
		int freq = 0, freq_sum = 0, t_freq = 0;
		for(int j = 0; j < elem.size(); ++ j)
		{
			t_label = elem[j];
			t_freq = (*q2f.find(t_label)).second;
			elem_info.push_back(pair<string, int>(t_label, t_freq));
			freq_sum += t_freq;
			t_score = lang_model->prob(t_label, lm);
			if(t_score > score)
			{
				score = t_score;
				label = t_label;
			}
		}
		label = raw_refine(label, lm, word2corr);
		map<string, vector<pair<string, int>>>::iterator itr = l2g.find(label);
		if(itr == l2g.end())
			l2g.insert(pair<string, vector<pair<string, int>>>(label, elem_info));
		else{
			for(int k = 0; k < elem.size(); ++ k)
				(*itr).second.push_back(elem_info[k]);
		}
	}
}

void DefaultModel::parse_info(string cluster, map<string, int> &q2f, string &cid)
{
	int del1 = cluster.find("\t"), del2 = cluster.find("\t", del1 + 1);
	
	cid = cluster.substr(del1 + 1, del2 - del1 - 1);
	string cluster_elem_s = cluster.substr(del2 + 1, cluster.length() - del2 - 1);
	vector<string> qf_list = util::tokenize_h(cluster_elem_s, " | ");
	//compute the language model
	for(int i = 0; i < qf_list.size(); ++ i)
	{
		string q_f = qf_list[i];
		int del = q_f.find("\t");
		string query = q_f.substr(0, del);
		string freq_c = q_f.substr(del + 1, q_f.length() - del - 1);
		int freq = atoi(freq_c.c_str());
		string q = util::format(query);
		//count query frequency
		map<string, int>::iterator q2f_itr = q2f.find(q);
		if(q2f_itr == q2f.end())
			q2f.insert(pair<string, int>(q, freq));
		else
			(*q2f_itr).second += freq;
		
	}
}

void DefaultModel::gen_language_model(map<string, int> &q2f, map<string, int> &lm)
{
	map<string, int>::iterator itr;
	for(itr = q2f.begin(); itr != q2f.end(); ++ itr)
	{
		string query = (*itr).first;
		int freq = (*itr).second;
		//generate language model
		co_occur(query, freq, lm);
	}
}
	
string DefaultModel::short_cut(map<string, vector<pair<string, int>>> &l2g)
{
	stringstream ss;
	map<string, vector<pair<string, int>>>::iterator itr;
	for(itr = label2group.begin(); itr != label2group.end(); ++ itr)
	{
		int total = 0;
		string lab = (*itr).first;
		vector<pair<string, int>> v = (*itr).second;
		for(int i = 0; i < v.size(); ++ i)
			total += v[i].second;
		ss << lab << "\t" << total << ":";
		for(int i = 0; i < v.size(); ++ i)
			ss << v[i].first  << "|";
		ss << endl;
	}	
	return ss.str();
}
	
void DefaultModel::build(string cluster)
{
	clear();
	parse_info(cluster, query2freq, cluster_id);
	//cout << "parse_info done\n";
	gen_language_model(query2freq, frequency);
	//cout << "gen_language_model done\n";	
	raw_grouping(query2freq, groups, dist_calc);
	//cout << "raw_grouping done\n";	
	elect_label(groups, query2freq, frequency, label2group);	
	//cout << "elect_label done\n";
	trans_model->build(label2group, translate_freq);
	//cout << "trans_model done\n";
	corr_model->build(frequency, translate_freq);
	//cout << "corr_model->build done\n";
	corr_model->correct_label(label2group, c_label);
	//cout << "corr_model->correct_label done\n";
	abbr_model->build(c_label, abbr_freq);
	//cout << "abbr_model->build done\n";
}

map<string, int> DefaultModel::getTransModel()
{
	return translate_freq;
}

map<string, int> DefaultModel::getLangModel()
{
	return frequency;
}

map<string, int> DefaultModel::getLogSortCut()
{
	return query2freq;
}

map<string, int> DefaultModel::getAbbrModel()
{
	return abbr_freq;
}

void DefaultModel::clear()
{
	cluster_id = "";
	translate_freq.clear();
	groups.clear();
	label2group.clear();
	frequency.clear();
	query2freq.clear();
	abbr_freq.clear();
	c_label.clear();
}